function [Uavg,Vavg,tausq_v_samples,rmse_samples] = run_modQTL_confounder(Y, J, burnin, ngibbs)
%
% [Uavg,Vavg,tausq_v_samples,rmse_samples] = modQTL_confounder(Y, J, burnin, ngibbs)
%
% confounder correction by sparse matrix factorization
% using blocked Gibbs sampling
%
% Y = U V
%
% code: Yongjin Park, ypp@csail.mit.edu
%

    [nprobe,nind] = size(Y);
    SIGSQ0 = var(Y(:));
    sigsq = SIGSQ0;

    U = 0.1 * randn(nprobe,J,'single');
    Vt = zeros(nind,J,'single');
    invtausq_u = 0.01*nind;
    invtausq_v = 0.01*nprobe*ones(J,1,'single');

    tausq_v_samples = NaN(J,ngibbs,'single');
    rmse_samples = NaN(100+burnin+ngibbs,1,'single');

    Uavg = zeros(size(U),'single');
    Vavg = zeros(size(Vt'),'single');

    INIT_TIME = 100;
    vsq = zeros(1,J,'single');

    for iter = 1:(INIT_TIME + burnin + ngibbs),

        % ================ sample V ================
        % (Y - U*Vt' + uu*vvt')' * uu;
        % (Y' - Vt*U' + vvt*uu') * uu;
        % (Y'*uu - Vt*U'*uu + vvt*uu'*uu)

        YtU = Y'*U;
        UtU = U'*U;
        VtUtU = Vt*UtU;

        for jj = 1:J,
            vvt_old = Vt(:,jj);
            denom_v = UtU(jj,jj) + invtausq_v(jj);

            vvt = (YtU(:,jj) - VtUtU(:,jj) + vvt_old*UtU(jj,jj))/denom_v;
            vsq(jj) = sum(vvt.^2);
            vvt = vvt + randn(nind,1,'single') * sqrt(sigsq/denom_v);

            delt_vvt = vvt - vvt_old;
            VtUtU = VtUtU + bsxfun(@times, delt_vvt, UtU(:,jj).');
            Vt(:,jj) = vvt;
        end


        if iter > INIT_TIME,
            invtausq_v = max( 0.1/nind, min( nind*10, sigsq * nind ./ vsq ) );
        end
        

        % ================ sample U ================
        % (Y - U*Vt' + u_j*v_j')*v_j
        % Y*v_j - U*Vt'*v_j + u_j*v_j'*v_j
        % Y*Vt_j - U*(V*Vt)_j + (V*Vt)_jj * u_j

        YVt = Y*Vt;
        VVt = Vt'*Vt;
        UVVt = U*VVt;

        for jj = 1:J,
            uu_old = U(:,jj);
            denom_u = VVt(jj,jj) + invtausq_u;

            uu = (YVt(:,jj) - UVVt(:,jj) + uu_old * VVt(jj,jj))/denom_u;
            uu = uu + randn(nprobe,1,'single')*sqrt(sigsq/denom_u);
            
            delt_uu = uu - uu_old;
            UVVt = UVVt + bsxfun(@times, delt_uu, VVt(:,jj).');
            U(:,jj) = uu;
        end
        
        R = Y - U*Vt';
        rss = sum(R(:).^2);
        nobs = nprobe*nind;

        rmse = sqrt(rss / nobs);
        rmse_samples(iter) = rmse;

        sigsq = (rss + 0.1*SIGSQ0*(J*nprobe + J*nind));
        sigsq = sigsq / (nobs + J*nprobe + J*nind);

        if iter > (INIT_TIME+burnin),
            tausq_v_samples(:,iter - INIT_TIME - burnin) = 1./invtausq_v;
            Uavg = Uavg + U / ngibbs;
            Vavg = Vavg + Vt' / ngibbs;
        end

        fprintf(2,'Iter = %05d, rmse = %.4f, sigma = %.4f\n', ...
                iter, rmse, sqrt(sigsq));

    end

end % end of function
